function [phi,psi,J,x] = sn(in)
% function [phi,psi] = sn(in)
%   This function solves the one-speed transport equation in slab geometry
%   using the discrete ordinates method.  Both fixed source and eigenvalue
%   problems can be performed.  Only isotropic sources and scattering are
%   considered.
%  
%   Inputs: 
%       in.numg     -- number of groups
%       in.numm     -- number of materials
%       in.xcm      -- slab bounds
%       in.xfm      -- number of fine meshes
%       in.mt       -- slab material ids
%       in.data     -- cross-sections in the form
%                      (mat1,g1) sT sA  nuSF chi sSg1->g1  sSg1->g2 ...
%                        ...,g2) sT sA  nuSF chi sSg2->g1  ....
%       in.src      -- volume source [1/cc-s]
%       in.ord      -- number of ordinates
%       in.ptype    -- 0 = fixed source, 1 = eigenvalue
%       in.maxinit  -- max inner iterations
%       in.maxinerr -- max pointwise phi error for inner iterations
%       in.maxoutit -- max outer iterations
%       in.maxouterr-- max pointwise fission source error for outers
%       in.maxkerr  -- max error for keff
%       in.bcL      -- left boundary condition
%       in.bcR      -- right boundary condition
%       in.acc      -- acceleration (0=none,1=cmr,2=cmfd,3=cmsc)
%   Output
%       phi         -- fine mesh-centered scalar flux
%       psi         -- fine mesh-edge angular flux
%       khist       -- keff history (if applicable)
%       shist       -- source history (if applicable)
%
% TO DO:
%       - implement coarse mesh acceleration for eigenvalue problems
%       - add multigroup capability to coarse mesh routines
%       - add reflective conditions (this should help in investigating
%         spectral radius of cmsc)

% assign angles and weights
[mu,wt] = S_1D(in.ord);
% allocate
nfm    = sum(in.xfm);      % total number of fine meshes
ncm    = length(in.xfm);   % number of coarse meshes
S      = zeros(nfm,in.ord,in.numg);     % fine mesh source external
fmmid  = zeros(nfm,1);     % fine mesh material id
% compute discretization
j = 0;
for i = 1:ncm
    dx((j+1):(j+in.xfm(i)))      = ( in.xcm(i+1) - in.xcm(i) ) / in.xfm(i);
    for g = 1:in.numg
        S((j+1):(j+in.xfm(i)), :, g)    = 0.5*in.src(g,i); % uniform, iso.
    end
    fmmid( (j+1):(j+in.xfm(i)))  = in.mt(i);  
    j = sum(in.xfm(1:i));
end
x = zeros(nfm,1);
x(1)=dx(1)/2;
for i = 2:length(dx)
    x(i) = x(i-1) + 0.5*(dx(i-1)+dx(i));
end
% precompute coefficients, following Eqs. 20.16-20.20 in 22.106 notes.
alpha = zeros(in.ord,1); % alpha=0 --> diamond difference; could add more
A     = zeros(nfm,4,in.numg);
B     = A;
for i = 1:nfm
    m = fmmid(i); 
    for n = 1:in.ord % use the sign of mu to choose appropriate form
        for g = 1:in.numg
            smu      = sign(mu(n)); 
            SigT     = in.data((m-1)*in.numg+g,1);
            denom    = 2*mu(n)+smu*(1+smu*alpha(n))*SigT*dx(i);
            A(i,n,g) = (2*mu(n)-smu*(1-smu*alpha(n))*SigT*dx(i)) /denom;
            B(i,n,g) = smu * 2 * dx(i) / denom;
        end
    end
end
psi    = zeros(nfm+1,in.ord,in.numg);   % current angular flux 
phi    = zeros(nfm,in.numg);            % current scalar flux
% solve the problem with the right solver
if in.ptype == 0
    [phi psi] = fixedsource(phi,psi,in,alpha,dx,fmmid,S,A,B,mu,wt,nfm,ncm);
    J = zeros(length(phi),1);
    for i = 1:length(mu)
        J = J + 0.5*(psi(1:end-1,i,1)+psi(2:end,i,1))*mu(i)*wt(i);
    end
    
elseif in.ptype == 1
    [phi,psi] = eigenvalue(phi,psi,in,alpha,dx,fmmid,S,A,B,mu,wt,nfm,ncm);
else
    disp('wrong ptype')
end

end
% end function sn

function [phi psi] = eigenvalue(phi,psi,in,alpha,dx,fmmid,S,A,B,mu,wt,nfm,ncm)
% function [phi psi] = eigenvalue(in,alpha,dx,fmmid,S,A,B,mu,wt,nfm,ncm)
%   This function solves an eigenvalue problem.  The algorithm is 
%
%   guess Sf and keff
%   while Sf and keff not converged
%       solve T psi^{m} = Sf^{m-1}/keff^{m-1}
%       update Sf and keff
%   end
%
%   The fixed source solver is used to the solve the inner equation, where
%   the source S is set to the current fission source.
fd  = zeros(nfm,1);
phi = ones(nfm,in.numg);
for i = 1:nfm
    m = fmmid(i);
    for g = 1:in.numg
        fd(i) = 0;
        %phi(i) = in.src(1,m);
        for gg = 1:in.numg
            fd(i) = fd(i) + phi(i,gg)*in.data((m-1)*in.numg+gg,3);
        end
        S(i,:,g) = in.data((m-1)*in.numg+g,4) * fd(i) * 0.5;
    end
end
for i = 1:nfm
   for g = 1:in.numg
      fd(i) = fd(i) + sum( S(i,:,g)'.*wt(:) );
   end
end
nrmfd   = sum(dx.*fd');
fd      = fd/nrmfd;
S       = S / nrmfd; % starts iff as our fission source
phi     = phi / nrmfd;
err_k   = 1.0;
err_s   = 1.0;
it      = 0;
k1      = 1.0;

% begin outers
while ( (err_k > in.maxkerr || err_s > in.maxouterr) && it < in.maxoutit ) 
    fd0 = fd;
    k0 = k1;
    if in.acc_e == 1
        % perform one outer iteration by cmfd
        [phi psi k1 S] = cmfd_eigen(phi,psi,in,dx,fmmid,S,fd,k1,mu,wt,nfm,ncm);
        
        % return
        for i = 1:nfm
            m = fmmid(i);
            for g = 1:in.numg
                fd(i) = 0;
                for gg = 1:in.numg
                    fd(i) = fd(i) + phi(i,gg)*in.data((m-1)*in.numg+gg,3);
                end
                S(i,:,g) = in.data((m-1)*in.numg+g,4) * fd(i) * 0.5;
            end
        end
    end
    [phi psi] = fixedsource(phi,psi,in,alpha,dx,fmmid,S/k1,A,B,mu,wt,nfm,ncm);
    % save old fission density
   % fd0 = fd;
   % k0 = k1; 
    % compute new fission density and source S
    for i = 1:nfm
       m = fmmid(i);
       for g = 1:in.numg
          fd(i) = 0;
          for gg = 1:in.numg
            fd(i) = fd(i) + phi(i,gg)*in.data((m-1)*in.numg+gg,3);
          end  
          S(i,:,g) = in.data((m-1)*in.numg+g,4) * fd(i) * 0.5;
       end
    end  
    k1 = k0*sum(dx'.*fd)/sum(dx'.*fd0);
    err_k = abs(k1-k0);
    err_s = max( abs(fd-fd0)./fd0 );
    it = it + 1;
    disp( ['OUTER ITERATION = ',num2str(it),'   k = ',num2str(k1),...
           '   err_k = ',num2str(err_k),'   err_s = ',num2str(err_s)] )
%     hold on
%     plot(fd,'Color',[rand rand rand])
%     pause()
end
    sprintf('   FINAL keff = %2.8f in %i iterations',k1,it )

end
% end function eigenvalue


function [phi psi] = fixedsource(phi,psi,in,alpha,dx,fmmid,S,A,B,mu,wt,nfm,ncm)
% function [phi psi] = fixedsource(phi,psi,in,alpha,dx,fmmid,S,A,B,mu,wt,nfm,ncm)
%   This function solves a fixed source problem defined by the mesh dx,
%   source S, and various constants A, B, mu, and wt.
Q      = zeros(nfm,in.ord,in.numg);     % emission density 
% group iterations, assuming no upscatter
for g = 1:in.numg
    % inner iteration convergence parameters
    err_phi = 1;    
    it      = 0; 
    while (err_phi > in.maxinerr && it <= in.maxinit )
        % Save old scalar flux
        phi0(:,1) = phi(:,g); 
        % Update emission density
        for i = 1:nfm
            % reset to fixed source (external or fission)
            Q(i,:,g) = S(i,:,g); 
            m = fmmid(i);
            for gg = 1:g
                % add scattering from groups gg>=g (only downscatter)
                Q(i,:,g) = Q(i,:,g) + ...
                           0.5*in.data((m-1)*in.numg+gg,4+g)*phi(i,gg);
            end
        end
        if in.bcL == 1
           psi(1,in.ord/2+1:end) = psi(1,in.ord/2:-1:1);
        end
        if in.bcR == 1
            psi(end,1:in.ord/2) = psi(1,in.ord:-1:in.ord/2+1);
        end
        
        [phi psi] = sweep(psi,phi,Q,alpha,A,B,mu,wt,nfm,g);
        % acceleration schemes (update both psi and phi)
        if ( in.acc == 1 )
            [tmp,phi] = cmr(psi,phi,S,in,mu,wt,dx,g);
        elseif ( in.acc == 2 )
            [tmp,phi] = cmfd(psi,phi,S,in,mu,wt,dx,g);
        elseif ( in.acc == 3)
            [tmp,phi] = cmsc(psi,phi,S,in,mu,wt,dx,g);
        elseif ( in.acc == 4)
            [tmp,phi] = cmsc2(psi,phi,S,in,mu,wt,dx,g);
        elseif ( in.acc == 5)
            [tmp,phi] = cmsc3(psi,phi,S,in,mu,wt,dx,g);              
        end      
        % Update error and iteration counter
        err_phi =  max(  abs(phi(:,g)-phi0)./phi0 );
        it      = it + 1;
      %  plot(1:length(phi(:,g)),phi(:,g),'c','LineWidth',2)
        if (mod(it,100)==0)
            disp(['      group = ',num2str(g),' it = ',num2str(it),'   err_phi = ',num2str(err_phi)])
        end        
    end
    if (it <= in.maxinit)
        disp(['   Group ',num2str(g),' converged in ', ...
              num2str(it),' iterations.'])  
    else
        disp(['   Group ',num2str(g),' failed to converge.'])  
    end
end  
end
% end function fixed source

function [phi,psi] = sweep(psi,phi,Q,alpha,A,B,mu,wt,nfm,g)
a1     = 1;         % angle indexes
b1     = length(mu)/2;  % negative from a1:b1, positive a2:b2
a2     = b1+1;
b2     = length(mu);
% Perform sweeps
for i = nfm:-1:1  % right-to-left
    psi(i,  a1:b1,g) = A(i,a1:b1,g).*psi(i+1,a1:b1,g) + ...
                       B(i,a1:b1,g).*Q(i,a1:b1,g);
end
for i = 1:nfm     % left-to-right
    psi(i+1,a2:b2,g) = A(i,a2:b2,g).*psi(i,a2:b2,g)   + ...
                       B(i,a2:b2,g).*Q(i,a2:b2,g);
end
% Update phi (need cell-centered psi, so use Eq. 20.15 of 106 notes)
for i = 1:nfm
    phi(i,g) = sum(wt(:).*(0.5*( (1+alpha(:)).*psi(i+1,:,g)'+ ...
        (1-alpha(:)).*psi(i,:,g)'   ) ));
end
end
% end function sweep


